import os
from PIL import Image
from tqdm import tqdm
import multiprocessing as mp
from functools import partial

def convert_to_gray(args):
    src_path, dst_path = args
    try:
        # 打开图像
        with Image.open(src_path) as img:
            # 转换为灰度图
            gray_img = img.convert('L').convert('RGB')
            # 确保目标文件夹存在
            os.makedirs(os.path.dirname(dst_path), exist_ok=True)
            # 保存灰度图
            gray_img.save(dst_path)
    except Exception as e:
        print(f"Error processing {src_path}: {str(e)}")

def process_imagenet(src_root, dst_root):
    """
    将ImageNet数据集转换为灰度图
    
    Args:
        src_root: 源ImageNet数据集路径
        dst_root: 目标保存路径
    """
    # 收集所有需要处理的图像路径
    tasks = []
    for class_folder in os.listdir(src_root):
        class_path = os.path.join(src_root, class_folder)
        if not os.path.isdir(class_path):
            continue
            
        for img_name in os.listdir(class_path):
            if not img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
                continue
                
            src_path = os.path.join(class_path, img_name)
            dst_path = os.path.join(dst_root, class_folder, img_name)
            tasks.append((src_path, dst_path))
    
    # 使用多进程处理
    print(f"Found {len(tasks)} images to process")
    with mp.Pool(processes=mp.cpu_count()) as pool:
        list(tqdm(pool.imap(convert_to_gray, tasks), total=len(tasks)))

if __name__ == '__main__':
    # 设置源数据集路径和目标路径
    src_root = '/ifs/data/PlatforMax/common/Dataset/ImageNet/ImageNet2012/train'  # 替换为您的ImageNet训练集路径
    dst_root = '/ifs/root/ipa01/101/user_101002/Datasets/ImageNet2012/train_gray'  # 替换为您想保存灰度图的路径
    
    process_imagenet(src_root, dst_root)